from src import Preproduction as Pre

id, ch = Pre.load("results/hanzi_2_one_hot.data")
cnt = len(ch)
V = 10000

print("cnt = " + str(cnt))
print("V = " + str(V))

from src.EmbeddingGloVe import genX as genX
from src.EmbeddingGloVe import Glove as Glove

pathBase = 'dataset/chat_corpus/clean_chat_corpus/'
paths = [
    "chatterbot.tsv",
    "douban_single_turn.tsv",
    "ptt.tsv",
    "qingyun.tsv",
    "subtitle.tsv",
    "tieba.tsv",
    "weibo.tsv"
]
fullPaths = [pathBase + path for path in paths]

import numpy as np

glove = Glove(V, 256)
glove.load("results/embedding120.model")

#i = glove.eval(id[c])

import random

# 从文档当中，采样batch组问答
# X = 问 Y = 答
# \t \n 视为END（id=0）
def Sample(rows, batch, seqLen, visualize=False):
    X_ = []
    Y_ = []
    
    for row in random.sample(rows, batch):
        s = row.split()
        x = s[0][-seqLen + 1:] + '\t'
        y = s[1]

        if visualize:
            print(x + " > " + y)

        X_.append(x)
        Y_.append(y)

    return (X_, Y_)

# 字符串转LSTM序列
def Transfer(x, seqLen):
    x = x[-seqLen + 1 :] + '\t'
    x = [glove.eval(0).detach().numpy()] * max(0, seqLen - len(x)) + \
        [glove.eval(1 if i == '\t' or i == '\n' else id[i] if id[i] < V else 0).detach().numpy() for i in x]
    return np.array([x])


def SampleChar(x, mode='top'):
    if mode == 'top':
        return np.argmax(x, axis=1)
    elif mode == 'rand':
        res = []
        for i in range(x.shape[0]):
            t = x[i].argsort()[-5:]
            p = [x[i][j] for j in t]
            c = t[np.random.choice(5, 1, p)[0]]
            res.append(c)
            print(ch[c])
        res = np.array(res, dtype=np.int32)
        return res
    return np.zeros(x.shape[0], dtype=np.int32)

from src.GeneratorLSTM import LSTM

embDim = 256
hidDim = 256
seqLen = 20

lstm = LSTM(embDim, hidDim, seqLen, V)
lstm.load("results/V3LSTM005.model")
lstm.evalMode()

def Gen(cur):
    Y = ''
    while True:
        X = Transfer(cur, seqLen)
        res = lstm.eval(X)
        c = SampleChar(res, 'rand')[0]
        if (c == 1):
            break
        else:
            c = ch[c] if c else ' '
            Y += c
            cur = (cur[1:] if len(cur) >= seqLen else cur) + c
    return Y

# 测试
'''
for i in range(1, 2):
    print("Epoch %d." % i)
    for j in range(len(paths)):
        print(paths[j])
        with open(fullPaths[j], "r") as f:
            rows = f.readlines()
            X_, Y_ = Sample(rows, 2, seqLen)
            for i in range(2):
                print(" Query> " + X_[i])
                print("Answer> " + Y_[i])
                print("   Gen> ", end='')
                cur = X_[i]
                Y = ''
                while True:
                    X = Transfer(cur)
                    res = lstm.eval(X)
                    c = SampleChar(res, 'rand')[0]
                    if (c == 0):
                        print('')
                        break
                    else:
                        Y += ch[c]
                        cur = (cur[1:] if len(cur) >= seqLen else cur) + ch[c]
                        print(ch[c], end='')
        print("------------------------------")
    print("==============================\n")
'''

while True:
    cur = input("Query> ")
    print("  Gen> " + Gen(cur))
    